#pragma once
#include <utility>

namespace lw
{
	enum Color
	{
		RED,
		BLACK
	};
	template<class K, class V>
	struct RBTreeNode
	{
		RBTreeNode<K, V>* _left;
		RBTreeNode<K, V>* _right;
		RBTreeNode<K, V>* _parent;
		Color _col;
		pair<K, V> _kv;
		RBTreeNode(const pair<K, V>& kv)
			:_left(nullptr)
			,_right(nullptr)
			,_parent(nullptr)
			,_col(RED)
			,_kv(kv)
		{}
	};
	template<class K, class V>
	class RBTree
	{
		typedef RBTreeNode<K, V> Node;
	public:
		bool Insert(const pair<K, V>& kv)
		{
			if (_root == nullptr)
			{
				_root = new Node(kv);
				_root->_col = BLACK;
				return true;
			}
			Node* curr = _root;
			Node* parent = nullptr;
			while (curr)
			{
				if (curr->_kv.first < kv.first)
				{
					parent = curr;
					curr = curr->_right;
				}
				else if (curr->_kv.first > kv.first)
				{
					parent = curr;
					curr = curr->_left;
				}
				else
				{
					return false;
				}
			}
			curr = new Node(kv);
			if (parent->_kv.first < kv.first)
				parent->_right = curr;
			else
				parent->_left = curr;
			curr->_parent = parent;

			while (parent && parent->_col == RED)
			{
				Node* grandfather = parent->_parent;
				if (parent == grandfather->_left)
				{
					Node* uncle = grandfather->_right;
					if (uncle && uncle->_col == RED)
					{
						parent->_col = uncle->_col = BLACK;
						grandfather->_col = RED;

						curr = grandfather;
						parent = curr->_parent;
					}
					else
					{
						if (curr == parent->_left)
						{
							//      g
							//   p     u
							//c
							RotatoR(grandfather);
							parent->_col = BLACK;
							grandfather->_col = RED;
						}
						else
						{
							//      g
							//   p     u
							//    c
							RotatoL(parent);
							RotatoR(grandfather);
							curr->_col = BLACK;
							grandfather->_col = RED;
						}
						break;
					}
				}
				else
				{
					Node* uncle = grandfather->_left;
					if (uncle && uncle->_col == RED)
					{
						parent->_col = uncle->_col = BLACK;
						grandfather->_col = RED;

						curr = grandfather;
						parent = curr->_parent;
					}
					else
					{
						if (curr == parent->_right)
						{
							//      g   
							//   u     p
							//           c
							RotatoL(grandfather);
							parent->_col = BLACK;
							grandfather->_col = RED;
						}
						else
						{
							//      g   
							//   u     p
							//        c
							RotatoR(parent);
							RotatoL(grandfather);
							curr->_col = BLACK;
							grandfather->_col = RED;
						}
						break;
					}
				}
				
			}
			_root->_col = BLACK;
			return true;
		}
		void RotatoL(Node* parent)
		{
			Node* subR = parent->_right;
			Node* subRL = subR->_left;
			parent->_right = subRL;
			if (subRL)
				subRL->_parent = parent;
			subR->_left = parent;
			Node* ppnode = parent->_parent;
			parent->_parent = subR;
			if (parent == _root)
			{
				_root = subR;
				subR->_parent = nullptr;
			}
			else
			{
				if (ppnode->_left == parent)
					ppnode->_left = subR;
				else
					ppnode->_right = subR;
				subR->_parent = ppnode;
			}
		}
		void RotatoR(Node* parent)
		{
			Node* subL = parent->_left;
			Node* subLR = subL->_right;
			parent->_left = subLR;
			if (subLR)
				subLR->_parent = parent;
			subL->_right = parent;
			Node* ppnode = parent->_parent;
			parent->_parent = subL;
			if (parent == _root)
			{
				_root = subL;
				subL->_parent = nullptr;
			}
			else
			{
				if (ppnode->_left == parent)
					ppnode->_left = subL;
				else
					ppnode->_right = subL;
				subL->_parent = ppnode;
			}
		}
		void InOrder()
		{
			_InOrder(_root);
		}
		bool IsBalance()
		{
			if (_root && _root->_col == RED)
				return false;
			Node* left = _root;
			int count = 0;
			while (left)
			{
				if (left->_col == BLACK)
					count++;
				left = left->_left;
			}
			return check(_root, 0, count);
		}
	private:
		bool check(Node* root, int count, int refBlackNumber)
		{
			if (root == nullptr)
			{
				if (count == refBlackNumber)
					return true;
				else
					return false;
			}
			if (root->_col == RED && root->_parent->_col == RED)
				return false;
			if (root->_col == BLACK)
				count++;
			return check(root->_left, count, refBlackNumber)
				&& check(root->_right, count, refBlackNumber);
		}
		void _InOrder(Node* root)
		{
			if (root == nullptr)
				return;
			_InOrder(root->_left);
			cout << root->_kv.first << " : " << root->_kv.second << endl;
			_InOrder(root->_right);
		}
		Node* _root = nullptr;
	};
}